// Assumptions: Supports only one client at a time
// server == BOB

#include "ramprakash.Part2Common.h"

typedef struct _user_record 
{
    char *u;
    char *p;
} user_t;


// Globals
struct sockaddr_in clnt;
int listenSockFd;
int clntSockFd;
fd_set readFds;
BIGNUM *gP;
BIGNUM *gGXA;
BIGNUM *gKAS;
BIGNUM *gXB;
BIGNUM *gNB; // Bob's Nonce
BIGNUM *gNA; // Alice's Nonce
char gDesKey[8];

// User records - add to this to add new users - Need to store this encrypted
// to make it more secure - FIXME fast!!
static user_t users[] = { 
    {"guest", "password"},
    {"user", "password"},
    {"admin", "admin"},
    {"ee282", "nice282"},
    {NULL, NULL}
};
                 

//socket helper functions
int createListeningSocket()
{
    struct sockaddr_in servaddr;
    int optval = 1;
    int buflen = 2048;

    listenSockFd = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listenSockFd < 0) 
    {
        perror("socket");
        return errno;
    }

    setsockopt(listenSockFd, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(int));
    setsockopt(listenSockFd, SOL_SOCKET, SO_SNDBUF, &buflen, sizeof(int));
    setsockopt(listenSockFd, SOL_SOCKET, SO_RCVBUF, &buflen, sizeof(int));


    memset(&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_addr.s_addr = htonl(LISTENING_ADDR);
    servaddr.sin_port = htons(LISTENING_PORT);

    if (bind(listenSockFd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
    {
        perror("bind");
        return errno;
    }

    if (listen(listenSockFd, 5) < 0)
    {
        perror("listen");
        return errno;
    }

    return listenSockFd;
}

// File utils
int sendFile(char *file)
{
    msg_t msg;
    FILE  *fd;
    file_t f;
    struct stat fInfo;
    int deslen;
    int j;

    memset(&f, 0, sizeof(f));
    memset(&msg, 0, sizeof(msg_t));
    memset(&fInfo, 0, sizeof(struct stat));

    fd = fopen(file, "r");
    if (!fd)
    {
        perror("fopen");
        return errno;
    }

    if (fstat(fileno(fd), &fInfo))
    {
        perror("fstat");
        return errno;
    }

    msg.cmd = FILE_RESP;
    strcpy(f.file, file);
    f.size = fInfo.st_size;
    f.chunk_size = fread(f.data, 1, 1024, fd);
    while ((f.offset + f.chunk_size) <= f.size) {
        // In the last packet for the file, send the SHA-1 digeest also
        if (f.chunk_size < 1024) {
            sha1Digest(file, f.digest);
            LOG("The calculated digest : ");
            for (j = 0; j < 20; j++) {
                LOG("%02x", f.digest[j]);
            }
            LOG("\n");
#ifdef TESTING
            f.digest[16] = 0xca;
            f.digest[17] = 0xfe;
            f.digest[18] = 0xca;
            f.digest[19] = 0xfe;
#endif
        }
        deslen = desCryptM(gDesKey, (char *)&f, sizeof(f), (char *)&msg.u.f, 1);
        msg.len = deslen;
        if (sendMsg(clntSockFd, (char *)&msg, sizeof(msg_t)) != 0) {
            return errno;
        }
        if (f.chunk_size < 1024) {
            if (feof(fd)) {
                LOG("Transmitted file: %s\n", file);
                break;
            }
        }
        f.offset += f.chunk_size;
        f.chunk_size = fread(f.data, 1, 1024, fd);
    }
    return 0;
}

// primitive authentication - can be improved.... a LOT (FIXME fast!!)
int authStatus (char *u, char *p)
{
    if (!p) {
        return -1;
    }
    int i = 0;
    while (users[i].u != NULL) {
        if ((strcmp(users[i].u, u) == 0) &&
            (strcmp(users[i].p, p) == 0)) {
            return 0;
        }
        i++;
    }
    return -1;
}

char *getPasswd(char *u)
{
    int i = 0;
    while (users[i].u != NULL) {
        if (strcmp(users[i].u, u) == 0) {
            return(users[i].p);
        }
        i++;
    }
    return NULL;
}

void generateAuthenticationParams()
{
    BN_CTX     *ctx;
    int         rand = 0;
    char        rand_buf[10];
#ifdef DEBUG
    char *bn_buf;
    int   llen;
#endif

    rand = getRandom();
    memset(rand_buf, 0, sizeof(rand_buf));
    sprintf(rand_buf, "%x", rand);

    // Bob's context:
    // Generate gXB - ensure it is greater than gP by generating random
    // number and adding it to gP
    if (BN_hex2bn(&gXB, rand_buf) == 0) {
        printf("BN_hex2bn failed\n");
        exit(1);
    }
    BN_add(gXB, gP, gXB);
#ifdef DEBUG
    llen = BN_num_bytes(gXB);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gXB);
    printf("gXB: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif

    //Calculate gGXA ^ gXB mod gP
    ctx = BN_CTX_new();
    gKAS  = BN_new();
    BN_mod_exp(gKAS, gGXA, gXB, gP, ctx);
#ifdef DEBUG
    llen = BN_num_bytes(gKAS);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gKAS);
    printf("gKAS: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif

    rand = getRandom();
    memset(rand_buf, 0, sizeof(rand_buf));
    sprintf(rand_buf, "%x", rand);
    if (BN_hex2bn(&gNB, rand_buf) == 0) {
        printf("BN_hex2bn failed\n");
        exit(1);
    }
    BN_add(gNB, gKAS, gNB);

#ifdef DEBUG
    llen = BN_num_bytes(gNB);
    bn_buf = malloc(llen);
    bn_buf = BN_bn2hex(gNB);
    printf("gNB: %s of length: %d\n", bn_buf, llen);
    free(bn_buf);
#endif
}

void sendMessage4()
{
    msg_t msg;
    int   len;
    int   deslen;
    char *buf;

    memset (&msg, 0, sizeof(msg));
    msg.cmd = MSG_4;

    //Encrypt using DES with gKAS(first 8 bytes) as key and 
    //encrypt the Nonce gNB
    len = BN_num_bytes(gNA);
    len *= 2;
    buf = BN_bn2hex(gNA);
    deslen = desCryptM(gDesKey, buf, len, msg.u.m4.enc_des_na, 1);
    msg.u.m4.len_des_na = deslen;

    sendMsg(clntSockFd, (char *)&msg, sizeof(msg));
}

void sendMessage2(char *pw)
{
    msg_t msg;
    int   len;
    int   deslen;
    char *buf;

    generateAuthenticationParams();

    memset (&msg, 0, sizeof(msg));
    msg.cmd = MSG_2;

    //Use our own encryption with password as the key and encrypt gKAS
    len = BN_num_bytes(gKAS);
    len *= 2;
    buf = BN_bn2hex(gKAS);
    memcpy(gDesKey, buf, 8);
    cryptM(pw, buf, len, msg.u.m2.enc_gkas_modp, 1);
    msg.u.m2.len_gkas_modp = len;

    //Encrypt using DES with gKAS(first 8 bytes) as key and 
    //encrypt the Nonce gNB
    len = BN_num_bytes(gNB);
    len *= 2;
    buf = BN_bn2hex(gNB);
    deslen = desCryptM(gDesKey, buf, len, msg.u.m2.enc_des_kas, 1);
    msg.u.m2.len_des_kas = deslen;

    sendMsg(clntSockFd, (char *)&msg, sizeof(msg));
}

int executeCmd(msg_t *msg)
{
    char u[32];
    char p[64];
    char pt[1500];
    char file[80];
#ifdef DEBUG
    char *bn_buf;
    int   llen;
#endif
    BIGNUM *bn = NULL;
    msg_t newMsg;
    memset(&newMsg, 0, sizeof(newMsg));
    switch (msg->cmd)
    {
        case MSG_1:
            if (authStatus(msg->u.m1.user, getPasswd(msg->u.m1.user)) != 0) {
                newMsg.cmd = AUTH_RESP;
                strcpy(newMsg.u.a.user, "Fail");
                if (sendMsg(clntSockFd, (char *)&newMsg, sizeof(msg_t)) != 0) {
                    return errno;
                }
                return 1;
            }

            memset(&pt, 0, sizeof(pt));
            cryptM(getPasswd(msg->u.m1.user), msg->u.m1.enc_p,
                             msg->u.m1.len_p, pt, 0);
            BN_hex2bn(&gP, pt);
#ifdef DEBUG
            llen = BN_num_bytes(gP);
            bn_buf = malloc(llen);
            bn_buf = BN_bn2hex(gP);
            printf("gP: %s of length: %d\n", bn_buf, llen);
            free(bn_buf);
#endif
            memset(&pt, 0, sizeof(pt));
            cryptM(getPasswd(msg->u.m1.user), msg->u.m1.enc_gxa_modp,
                             msg->u.m1.len_gxa_modp, pt, 0);
            BN_hex2bn(&gGXA, pt);

            sendMessage2(getPasswd(msg->u.m1.user));

            break;

        case MSG_3:
            memset(&pt, 0, sizeof(pt));
            desCryptM(gDesKey, msg->u.m3.enc_des_nb, 
                      msg->u.m3.len_des_nb, pt, 0);
            BN_hex2bn(&bn, pt);
            if (BN_cmp(bn, gNB) != 0) {
                printf("Invalid Nonce - Rejecting Authentication\n");
                newMsg.cmd = AUTH_RESP;
                strcpy(newMsg.u.a.user, "Fail");
                if (sendMsg(clntSockFd, (char *)&newMsg, sizeof(msg_t)) != 0) {
                    return errno;
                }
                return 1;
            }
            
            desCryptM(gDesKey, msg->u.m3.enc_des_na, 
                      msg->u.m3.len_des_na, pt, 0);
            BN_hex2bn(&gNA, pt);
#ifdef DEBUG
            llen = BN_num_bytes(gNA);
            bn_buf = malloc(llen);
            bn_buf = BN_bn2hex(gNA);
            printf("gNA: %s of length: %d\n\n", bn_buf, llen);
            free(bn_buf);
#endif
            sendMessage4();
            break;

        case AUTH_REQ:
            desCryptM(gDesKey, msg->u.a.user, 32, u, 0);
            desCryptM(gDesKey, msg->u.a.passwd, 64, p, 0);
            if (authStatus(u, p) == 0)
            {
                newMsg.cmd = AUTH_RESP;
                strcpy(newMsg.u.a.user, "Success");
            }
            else
            {
                newMsg.cmd = AUTH_RESP;
                strcpy(newMsg.u.a.user, "Fail");
            }

            if (sendMsg(clntSockFd, (char *)&newMsg, sizeof(msg_t)) != 0) {
                return errno;
            }
            break;

        case FILE_REQ:
            desCryptM(gDesKey, msg->u.f.file, 80, file, 0);
            if (sendFile(file) != 0) {
                return FILE_RESP;
            }
            break;

        default:
            printf("Unknown command\n");
            exit(-1);
            break;
    }

    return 0;
}

int dispatcher()
{
    msg_t msg;
    socklen_t len;
    while (1) 
    {
        memset(&msg, 0, sizeof(msg));
        // Blocking select
        select(clntSockFd + 1, &readFds, NULL, NULL, NULL);

        if ((len = recvfrom(clntSockFd, &msg, sizeof(msg), 0, NULL, 0)) > 0) {
            if (executeCmd(&msg) == FILE_RESP) {
                msg.cmd = FILE_RESP;
                msg.u.f.size = 0;
                sendMsg(clntSockFd, (char *)&msg, sizeof(msg_t));
            }
        }
        else {
            //perror("recvfrom");
            break;
        }
    }

    close(clntSockFd);
    return 0;
}

int waitForClient()
{
    int buflen = 2048;
    socklen_t len;
    clntSockFd = accept(listenSockFd, (struct sockaddr *)&clnt, &len);
    if (clntSockFd < 0) 
    {
        perror("accept");
        return -1;
    }

    setsockopt(clntSockFd, SOL_SOCKET, SO_SNDBUF, &buflen, sizeof(int));
    setsockopt(clntSockFd, SOL_SOCKET, SO_RCVBUF, &buflen, sizeof(int));

    FD_SET(clntSockFd, &readFds);
    return 0;
}

#ifdef DESDEBUG
int main(int argc, char *argv[])
#else
int main()
#endif
{
#ifdef DESDEBUG
    // This piece of code is just to test the desCryptM DES utility
    // function. run tcpSrvr passing any string to encrypt and then decrypt the
    // ciphertext back. To compile with this flag set, run "make desdebug"
    // If your input string has whitespace, enclose the string in double quotes
    char key[8] = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08};
    char *string;
    char outbuf[20];
    char inbuf[20];
    int deslen;
    int deslen2;

    if (argc != 2) {
        printf("Usage: %s <string to encrypt>\n", argv[0]);
        exit(1);
    }

    string = argv[1];
    deslen = desCryptM(key,  string, strlen(string), outbuf, 1);
    printf("Plaintext len(%lu): %s ==> Ciphertext len(%d): %s\n", 
            strlen(string), string, deslen, outbuf);
    deslen2 = desCryptM(key,  outbuf, deslen, inbuf, 0);
    printf("Ciphertext len(%d): %s ==> Plaintext len(%lu): %s\n", 
            deslen, outbuf, strlen(inbuf), inbuf);
    return 1;
#endif
    if (createListeningSocket() < 0) {

        LOG("Socket creation failed - Exiting\n");
        return -1;
    }

    while(1) {

        FD_ZERO(&readFds);

        if (waitForClient() < 0) {
            LOG("accept failed\n");
            return -1;
        }

        dispatcher();
    }
}
